package rawDataExtraction;

import java.awt.Color;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

import javax.imageio.ImageIO;

import org.omg.CORBA.ARG_OUT;

import projectUtils.ByteUtil;

public class MNIST_DigitToImages {
	public MNIST_DigitToImages()
	{
		
	}
	
	public static void ExtractDigitImages(String setType)
	{
		int numItems = 0;
		String labelName, dataName;
		String outputFolder;
		String userFolder = System.getProperty("user.dir");
		String dataFolder = userFolder + File.separator + "Training" + File.separator;
		if (setType.equals("Training"))
		{
			numItems = 60000;
			labelName = dataFolder + "train-labels.idx1-ubyte";
			dataName = dataFolder + "train-images.idx3-ubyte";
			outputFolder = dataFolder + "trainingSet" + File.separator;
		}
		else if (setType.equals("Testing"))
		{
			numItems = 10000;
			labelName = dataFolder + "t10k-labels.idx1-ubyte";
			dataName = dataFolder + "t10k-images.idx3-ubyte";
			outputFolder = dataFolder + "testSet" + File.separator;
		}
		else
			return;
		byte[] labelFile = null;
		byte[] dataFile = null;
		try {
			labelFile = Files.readAllBytes(Paths.get(labelName));
			dataFile = Files.readAllBytes(Paths.get(dataName));
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
		int[] fileCount = new int[10];
		for (int i = 0; i < 10; i++)
			fileCount[0] = 0;
		for (int i = 0; i < numItems; i++)
		{
			int labelIndex = 8 + i;
			int dataIndex = 8 + 4 + 4 + (28 * 28) * i;
			BufferedImage image = new BufferedImage(28, 28, BufferedImage.TYPE_INT_RGB);
			int label = ByteUtil.ByteToInt(labelFile[labelIndex]);
			String outputFile = outputFolder + label + File.separator + fileCount[label] + ".bmp";
			fileCount[label]++;
			for (int m = 0; m <= 27; m++)
			{
				for (int n = 0; n <= 27; n++)
				{
					int pixelIndex = dataIndex + m * 28 + n;
					int pixel = 255 - ByteUtil.ByteToInt(dataFile[pixelIndex]);
					image.setRGB(n, m, new Color(pixel, pixel, pixel).getRGB());
				}
			}
			try {
				ImageIO.write(image, "bmp", new File(outputFile));
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
	}
}
